package com.xpj.aspect;

import com.google.common.util.concurrent.RateLimiter;
import com.xpj.annotation.RequestLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

@Aspect
@Component
public class RateLimiterAspect {

    @Pointcut("@annotation(com.xpj.annotation.RequestLimiter)")
    public void pointcut(){}

    private ConcurrentHashMap<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>();

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint){
        try {
            MethodSignature signature = (MethodSignature)joinPoint.getSignature();
            RequestLimiter requestLimiter = signature.getMethod().getAnnotation(RequestLimiter.class);

            /**
             * 先获取限流策略，如果获取不到则代表第一次请求，需生成限流策略放入缓存中
             */
            RateLimiter rateLimiter = Optional.ofNullable(rateLimiters.get(requestLimiter.name())).orElseGet(() -> {
                RateLimiter limiter = RateLimiter.create(requestLimiter.token());
                /**
                 * 此处有可能多个线程都未获取到，会造成创建多个限流策略，
                 * 故使用putIfAbsent方式存放，如果返回值不为空，
                 * 则代表有线程已经放入的限流策略，则直接获取，否则返回刚创建的限流策略
                 */
                return Optional.ofNullable(rateLimiters.putIfAbsent(requestLimiter.name(), limiter)).orElse(limiter);
            });

            if (!rateLimiter.tryAcquire()){
                return "too many request";
            }

            Object result = joinPoint.proceed();

            return result;
        } catch (Throwable throwable) {
            return "error";
        }
    }

}
